# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import json
import re
import time
from typing import Dict
from events import scheduler
from exception import AppException
from common.cache.queue_cache import QueueCache
from common.cache.pool_cache import KeyPoolCache
from common.enums.agent import AgentEnum
from common.enums.models import ModelsEnum
from common.enums.public import CrontabEnum
from common.models.users import UserModel
from common.models.dev import DevModelsModel
from common.models.sys import SysCrontabModel
from common.models.agents import AgentKnowModel
from common.models.agents import AgentKnowArchiveModel
from common.models.agents import AgentKnowSplittingModel
from common.chain.chain_server import ChatUtils, AIChannelFactory


async def execute(**kwargs):
    """ 知识库QA拆分 """
    start_time = time.time()
    crontab_id: int = int(kwargs["w_id"])
    crontab_job: str = str(kwargs["w_job"])
    if not scheduler.get_job(crontab_job):
        scheduler.pause_job(crontab_job)

    try:
        kid = await QueueCache.pop_queue_value(QueueCache.QA_JOB)
        if kid is None:
            return True

        await handle_qa(int(kid))
        await SysCrontabModel.compute(crontab_id, start_time)
    except Exception as e:
        await SysCrontabModel.compute(crontab_id, start_time, status=CrontabEnum.CRON_ERROR, error=str(e))


async def handle_qa(kid: int):
    # 开始时间
    start_time = int(time.time())

    # 查询数据
    data = await AgentKnowSplittingModel.filter(id=kid).first()
    if not data:
        return True

    # 查询文件
    archive = await AgentKnowArchiveModel.filter(id=data.archive_id, is_delete=0).first()
    if not archive:
        return await update_data_error(data.archive_id, error="文档记录不存在")
    if archive.qa_status == AgentEnum.EMB_FAIL:
        return await update_data_error(data.archive_id, error=archive.fail_reason)
    if archive.qa_status == AgentEnum.EMB_WAIT:
        archive.qa_status = AgentEnum.EMB_ING
        archive.update_time = int(time.time())
        await archive.save()

    # 查知识库
    know = await AgentKnowModel.filter(id=data.know_id, is_delete=0).first()
    if not know:
        return await update_data_error(data.know_id, error="知识库已被删除")
    if know.is_disable:
        return await update_data_error(data.know_id, error="知识库已被禁用")

    # 验证模型
    try:
        models = await DevModelsModel.check_models(know.splits_model, ModelsEnum.TYPE_CHAT)
        chat_channel: str = models["channel"]
        chat_alias: str = models["alias"]
        chat_model: str = models["model"]
        chat_price: int = int(models["price"])
        chat_config: dict = models["config"]
    except Exception as e:
        return await update_data_error(archive.id, error=str(e))

    # 读取密钥
    try:
        apikey = await get_apikey(chat_channel, chat_config.get("check_key"))
    except Exception as e:
        return await update_data_error(archive.id, error=str(e))

    # 查询用户
    user = await UserModel.filter(id=know.user_id, is_delete=0).first()
    if not user:
        return await update_data_error(data.know_id, error="账号已注销")
    if user.is_disable == 1:
        return await update_data_error(data.know_id, error="账号已冻结")
    if user.points <= 0 < chat_price:
        return await update_data_error(data.know_id, error="账号余额不足")

    # 标记开始
    data.fail_reason = ""
    data.retry = data.retry + 1
    data.model_alias = chat_alias
    data.model_name = chat_model
    data.status = AgentEnum.EMB_ING
    data.update_time = int(time.time())
    await data.save()

    print(f"QA拆分任务: {kid}")
    try:
        # 发起请求
        answer = ""
        messages = await get_messages(data.content)
        aiServer = AIChannelFactory.create(chat_channel, chat_config, apikey)
        stream = await aiServer.llm(messages)
        async for chunk in stream:
            answer += chunk.choices[0].delta.content

        # 提取结果
        qaDict: dict = extract_qa_pairs(answer)
        usages: dict = ChatUtils.usage(messages, answer)
        points: int = ChatUtils.compute_price(usages["total_tokens"], chat_price)

        # 保存结果
        data.usage_points = points
        data.usage_tokens = json.dumps(usages, ensure_ascii=False)
        data.results = json.dumps(qaDict, ensure_ascii=False)
        data.status = AgentEnum.EMB_YES
        data.task_time = int(time.time()) - start_time
        data.update_time = int(time.time())
        await data.save()
    except Exception as e:
        print("QA拆分失败: " + str(e))
        data.status = AgentEnum.EMB_FAIL
        data.fail_reason = str(e)
        data.task_time = int(time.time()) - start_time
        data.update_time = int(time.time())
        await data.save()


async def get_apikey(channel: str, check_key: bool = True):
    """ 获取密钥 """
    apikey = {}
    if check_key:
        apikey = await KeyPoolCache(scene="chat", channel=channel).get_key() or {}
        if not apikey:
            raise AppException("当前模型尚未配置密钥,请与管理员联系")
    return apikey


async def get_messages(content: str):
    prompts = AgentEnum.get_prompt_tpl("qa")
    user = prompts["user"].replace("{{text}}", content)
    return [
        {"role": "system", "content": prompts["system"]},
        {"role": "user", "content": user}
    ]


async def update_data_error(fid: int, error: str):
    try:
        await AgentKnowArchiveModel.filter(id=fid).update(
            fail_reason=error,
            qa_status=AgentEnum.EMB_FAIL,
            update_time=int(time.time())
        )
    except Exception as e:
        print(str(e))


def extract_qa_pairs(text) -> Dict[str, str]:
    """
    从给定的文本中提取问题和答案对。

    Args:
        text (str): 包含问题和答案的文本。

    Returns:
        dict: 一个字典，其中键是问题内容，值是对应的答案内容。
    """
    pattern = r'Q\d+:\s*(.*?)\nA\d+:\s*(.*?)(?=\nQ\d+:|$)'
    matches = re.findall(pattern, text, re.DOTALL)
    return {question.strip(): answer.strip() for question, answer in matches}
